import gym
from gym import spaces
import numpy as np
from common_function import dcm_from_euler, rotation_matrix
from stable_baselines3 import DDPG
from stable_baselines3.common.noise import NormalActionNoise

class UAVDynamics:
    def __init__(self, params):
        self.params = params
        self.w_current = params.ModelInit_Rads

    def get_angular_velocity_derivative(self, M, J, omega):
        return np.linalg.inv(J) @ (M - np.cross(omega, J @ omega))

    def get_linear_velocity_derivative(self, F, m):
        return np.clip(F / m, -0.4, 0.4)

    def Obtain_force_torque(self, w, R, Cm, Ct, Vb, Cd, wb, Cdm, Jrp):
        M_rctcm = np.array([
            [-np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Roll torque
            [np.sin(np.pi / 4) * R * Ct, -np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Pitch torque
            [Cm, Cm, -Cm, -Cm]  # Yaw torque
        ])

        Mp = np.dot(M_rctcm, w ** 2)  # Torque

        Fp = np.array([0, 0, -np.sum(Ct * (w ** 2))])  # Thrust
        Fd = -Cd * Vb * np.abs(Vb) * 0.5  # Aerodynamic force
        Md = -Cdm * wb * np.abs(wb)  # Aerodynamic moment

        # Gyro moment
        Ga = np.zeros(3)
        Ga[0] = Jrp * wb[1] * (w[0] + w[1] - w[2] - w[3])
        Ga[1] = Jrp * wb[0] * (-w[0] - w[1] + w[2] + w[3])

        return Fp, Fd, Mp, Md, Ga

    def motor_dynamics(self, w_current, w_target, T, dt):
        dw = (w_target - w_current) / T
        return w_current + dw * dt

    def update_dynamics(self, state, w_target, dt):
        velB, angEuler, rateB = state

        m = self.params.ModelParam_uavMass
        J = self.params.ModelParam_uavJ
        g = self.params.ModelParam_envGravityAcc
        R = self.params.ModelParam_uavR
        Cm = self.params.ModelParam_rotorCm
        Ct = self.params.ModelParam_rotorCt
        Cd = self.params.ModelParam_uavCd
        Cdm = self.params.ModelParam_uavCCm
        Jrp = self.params.ModelParam_motorJm
        T = self.params.ModelParam_motorT

        self.w_current = self.motor_dynamics(self.w_current, w_target, T, dt)
        Fp, Fd, Mp, Md, Ga = self.Obtain_force_torque(self.w_current, R, Cm, Ct, velB, Cd, rateB, Cdm, Jrp)
        dcm = dcm_from_euler(angEuler)
        total_force = Fp + Fd + np.array([0, 0, m * g]).dot(dcm)
        total_torque = Mp + Md + Ga
        velB_dot = self.get_linear_velocity_derivative(total_force, m)
        rateB_dot = self.get_angular_velocity_derivative(total_torque, J, rateB)

        velB = np.clip(velB + velB_dot * dt, -3, 3)
        rateB = rateB + rateB_dot * dt
        angEuler = angEuler + rateB * dt

        return velB, angEuler, rateB

class UAVEnv(gym.Env):
    def __init__(self, params, waypoints):
        super(UAVEnv, self).__init__()
        self.params = params
        self.dynamics = UAVDynamics(params)
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(4,), dtype=np.float32)  # 4 motors
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(12,), dtype=np.float32)  # 12 states
        self.waypoints = waypoints
        self.current_waypoint_index = 0
        self.state = None
        self.done = False
        self.reset()

    def reset(self):
        self.state = (np.zeros(3), np.zeros(3), np.zeros(3))  # Reset to initial state: (velB, angEuler, rateB)
        self.done = False
        self.current_waypoint_index = 0
        return np.concatenate(self.state)

    def step(self, action):
        action = np.clip(action, self.action_space.low, self.action_space.high)
        velB, angEuler, rateB = self.state
        w_target = self.params.ModelInit_Rads + action * 1000  # Scale action to motor speeds
        next_state = self.dynamics.update_dynamics(self.state, w_target, 0.01)
        self.state = next_state
        reward = self.compute_reward(next_state)
        self.done = self.is_done(next_state)
        if self.done:
            self.current_waypoint_index += 1
            if self.current_waypoint_index < len(self.waypoints):
                self.done = False
        return np.concatenate(next_state), reward, self.done, {}

    def compute_reward(self, state):
        velB, angEuler, rateB = state
        target_pos = self.waypoints[self.current_waypoint_index]
        position_error = np.linalg.norm(velB[:3] - target_pos)
        reward = -position_error
        return reward

    def is_done(self, state):
        velB, angEuler, rateB = state
        target_pos = self.waypoints[self.current_waypoint_index]
        if np.linalg.norm(velB[:3] - target_pos) < 0.1:  # Close enough to target
            return True
        if np.linalg.norm(velB) > 10 or np.linalg.norm(angEuler) > np.pi or np.linalg.norm(rateB) > 10:
            return True
        return False

    def render(self, mode='human'):
        pass

    def close(self):
        pass
